import json
import time
from locust import HttpUser, task, between
import os
import glob

request_count = 0

def get_image_paths(root_dir):
    """
    获取指定目录及其所有子目录下的图片文件路径。

    参数:
        root_dir (str): 要遍历的根目录路径。

    返回:
        list: 包含匹配到的图片文件路径的列表。
    """
    # 支持的图片文件扩展名列表
    extensions = ['*.jpg', '*.jpeg', '*.png', '*.gif', '*.bmp']
    
    # 存储找到的图片路径
    image_paths = []

    # 遍历目录
    for ext in extensions:
        # os.path.join用于连接目录与文件名或目录；glob.glob用于获取匹配到的文件路径列表
        image_paths.extend(glob.glob(os.path.join(root_dir, '**', ext), recursive=True))

    return image_paths

root_directory = 'xun_ji_test/data/image/recognition/test_1'
image_list = get_image_paths(root_directory)

class QuickstartUser(HttpUser):
    wait_time = between(1, 2)

    # 图片路径的列表
    image_list = get_image_paths(root_directory)

    @task
    def upload_image(self):
        global request_count

        # 遍历图片列表，依次发送每个图片
        for image_path in self.image_list:
            print(image_path)
            start_time = time.time()
            with open(image_path, "rb") as image_file:
                response = self.client.post(
                    "https://iacc.gdcpzk.cn/api/recognize/",
                    files={"image": ("image.jpg", image_file, "image/jpeg")}
                )
            end_time = time.time()

            response_time = end_time - start_time
            request_count += 1
            print('Status code:', response.status_code)
            print(f"Request #{request_count} took {response_time} seconds")
            json_response = response.json()
            print(json_response)

            # 保存响应以备后用
            self.json_response = json_response
            # 将响应保存到json文件中
            with open('response_d.json', 'a') as f:
                json.dump(self.json_response, f)
                f.write('\n')


    